% QDMDipoleFittingBatch.m -- Main MATLAB script for fitting a magnetic dipole model
%    to Bz data using mixed linear and nonlinear least squares.
%
%  References: - R.R. Fu, E.A. Lima, M.W.R. Volk, R. Trubko (2020) "High-sensitivity
%  moment magnetometry with the quantum diamond microscope,"  Geochemistry,
%  Geophysics, Geosystems.
%              - E. A. Lima and B. P. Weiss (2016) "Ultra-high Sensitivity
%  Moment Magnetometry of Geological Samples Using Magnetic Microscopy,"
%  Geochemistry, Geophysics, Geosystems, 17, doi:10.1002/2016GC006487.
%
%  Remarks: inclination and declination follow the convention in the Fu et al. (2020) 
%    reference. The Bz map can be thought of as a map on Earth with N at 12 o'clock, 
%    E at 3 o'clock and up out of the page.  Inclination is positive if moment points into
%    the sample.
%
% -----------------------------------------------------------------------------------
%                Matlab code by Eduardo A. Lima and Roger R. Fu
%      Copyright (C) 2017-2020 MIT Paleomagnetism Lab, Harvard Paleomagnetics Lab
% -----------------------------------------------------------------------------------
function QDMDipoleFittingBatch(INFILE, XY, CROPFACT, DOWNSMPL, NRUNS, outputtrue,plottrue)

OutFileName='DipoleInversions.txt';     % Data file where moment inversion results will be saved

rng('shuffle')                          % Randomize seed for random number generator
set(0,'DefaultFigureColormap',jet)      % Comment this if you prefer parula color map


METHOD=0;       % 0 = non-linear least squares optimization , 1=Nelder-Mead (simplex) optimization
NOISE=0;        % 0 = no noise, 1 = add white noise to yield prescribed SNR
SNR=0;          % signal-to-noise ratio in dB (0 = 1:1, 20 = 10:1, etc.)

AUTO=0;         % automatically find dipole position from Bt map

MINTOL=1;       % number of points used in determining the optimal solution

DISPLAY=0;      % set this to 1 to visualize optimization process

m0=1e-11;       % rough estimate of order of magnitude of the moment to adjust converge criteria

step=0;
load(INFILE);
[filepath,name,ext]=fileparts(INFILE);

if DOWNSMPL % Downsample map by factor DOWNSMPL, if selected. This
    % helps speed up if map resolution is unnecessarily high
    Bz=downsample(downsample(Bz,DOWNSMPL)',DOWNSMPL)';
    step=step*DOWNSMPL;
end

scan=Bz;    % Bz is assumed in T

x=((1:size(scan,2))-1)*step;    % Create x-y coordinate grid
y=((1:size(scan,1))-1)*step;
[X,Y]=meshgrid(x,y);
figure(1);
if NOISE                        % Add white noise if applicable
    noise=sqrt(10^(-SNR/10)*var(scan(:)))*randn(size(scan));
    scan=scan+noise;
end

if DISPLAY                      % Only run code once when visualizing optimization
    NRUNS=1;
end

if DOWNSMPL                     % Adjust center of region of interest (cropping) if necessary
    j=round(XY(1)/DOWNSMPL);
    i=round(XY(2)/DOWNSMPL);
else
    j=round(XY(1));
    i=round(XY(2));
end

x0=j*step;
y0=i*step;

%Cropping window is (2*CROPFACT + 1) X (2*CROPFACT + 1) in size
cropj=round(1+j+[-CROPFACT CROPFACT]);
cropi=round(1+i+[-CROPFACT CROPFACT]);

% Adjust cropping area boundaries to fall whithin the map
scansize = size(scan);
for p=1:2
    if cropj(p) < 1
        cropj(p) = 1;
    end
    if cropi(p) < 1
        cropi(p) = 1;
    end
    if cropj(p) > scansize(2)
        cropj(p) = scansize(2);
    end
    if cropi(p) > scansize(1)
        cropi(p) = scansize(1);
    end
end

%crop field map
scanc=scan(cropi(1):cropi(2),cropj(1):cropj(2));
Xc=X(cropi(1):cropi(2),cropj(1):cropj(2));
Yc=Y(cropi(1):cropi(2),cropj(1):cropj(2));
xc=x(cropj(1):cropj(2));
yc=y(cropi(1):cropi(2));

%
imagesc(xc,yc,abs(Bz(cropi(1):cropi(2),cropj(1):cropj(2))));
axis xy, axis equal, axis tight
caxis([0 1]*max(abs(caxis)));
colormap(hot)
colorbar
title(sprintf('Cropped map'));
drawnow

P00(1)=x0;
P00(2)=y0;
P00(3)=h;

drawnow
P=zeros(length(P00)+3,NRUNS);
fval=zeros(1,NRUNS);
fval2=zeros(1,NRUNS);
for k=1:NRUNS
    if NRUNS==1
        P0=P00; % No randomization of origin coordinats guesses if running only once
    else
        P0=P00+0.2*(rand(size(P00))-0.5).*P00; % Introduce +/-10% perturbation to initial guesses for origin coordinates
    end
    options=optimset('TolX',10^(floor(log10(m0))-5),'TolFun', 10^(floor(log10(max(abs(Bz(:)))))-8),'MaxFunEvals',6000,'MaxIter',2000,'Display','none');
    
    if METHOD
        [P(1:3,k),fval2(k),exitflag,output] = fminsearch(@(Pp) SourceFitDipoleBz(Pp,Xc,Yc,scanc,DISPLAY,METHOD),P0,options); % Nelder-Mead optimization method
    else
        [P(1:3,k),fval2(k),resd,exitflag,output] = lsqnonlin(@(Pp) SourceFitDipoleBz(Pp,Xc,Yc,scanc,DISPLAY,METHOD),P0,[],[],options); %NLS method
    end
    [resid,BzModel,M]=SourceFitDipoleBz(P(1:3,k),Xc,Yc,scanc,DISPLAY,METHOD);   % Compute field map of solution when optimization is completed
    
    Mx=M(1);
    My=M(2);
    Mz=M(3);
    m=sqrt(Mx^2+My^2+Mz^2);
    theta=acosd(Mz/m);
    phi=atan2d(My,Mx);
    P(4,k)=m;
    %convert spherical angles to inclination and declination
    P(5,k)=90-theta;
    P(6,k)=phi-90;
    
    %enforce range for angular parameters
    change=1;
    while change
        change=0;
        if P(5,k)>90
            P(5,k)=180-P(5,k); % i' = 180-i
            P(6,k)=P(6,k)+180; % d' = d+180
            change=1;
        end
        if P(5,k)<-90
            P(5,k)=-180-P(5,k); % i' = -180-i
            P(6,k)=P(6,k)+180; % d' = d+180
            change=1;
        end
        if P(6,k)<0
            P(6,k)=P(6,k)+360; % d' = d+360
            change=1;
        end
        if P(6,k)>=360
            P(6,k)=P(6,k)-360; % d' = d-360
            change=1;
        end
    end
    if DISPLAY
        disp(sprintf('Moment: %.4e',P(4)))
        disp(sprintf('Inclination: %.1f',P(5)))
        disp(sprintf('Declination: %.1f',P(6)))
        disp(sprintf('Height: %.4e',P(3)))
    end
    
    
    fprintf('..%0d..(%0d)  ',k,output.iterations);
    if ~mod(k,10)
        fprintf('\n');
    end
    
    fval(k)=sqrt(sum(sum((BzModel-scanc).^2))/numel(scanc)); % compute normalized residual
end

% Pick best solution as the one with minimum cost (residual). MINTOL specifies
% whether this minimum is computed based on a single point or on an average of points with
% with minimal cost. In almost all most cases, MINTOL=1 should be used.

i0=find(fval==min(fval));
fsort=sort(fval);

i=find(fval<=fsort(MINTOL));
disp(sprintf('\n--- Averaging %d points ---',numel(i)))
if numel(i)>0.1*numel(fval)
    disp('Too many points are being averaged. Consider adjusting MINTOL parameter.')
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
Popt=zeros(size(P));
Popt(4)=sum(P(4,i).*fval(i))/sum(fval(i));
mopt=Popt(4)
disp(sprintf('(min = %1.3d)',P(4,i0)));

Popt(5)=sum(P(5,i).*fval(i))/sum(fval(i));
iopt=Popt(5);
-iopt
disp(sprintf('(min = %1.3f)',P(5,i0)));

Popt(6)=sum(P(6,i).*fval(i))/sum(fval(i));
dopt=Popt(6);                 
mod(360-dopt,360)
disp(sprintf('(min = %1.3f)',P(6,i0)));

Popt(1)=sum(P(1,i).*fval(i))/sum(fval(i));
Popt(2)=sum(P(2,i).*fval(i))/sum(fval(i));
Popt(3)=sum(P(3,i).*fval(i))/sum(fval(i));
hopt=Popt(3)
disp(sprintf('(min = %1.3d)',P(3,i0)));

xopt=Popt(1)
disp(sprintf('(min = %1.3d)',P(1,i0)));

yopt=Popt(2)
disp(sprintf('(min = %1.3d)',P(2,i0)));

% Display plots of cost function vs inversion parameter for all candidate solutions
% This helps assess how robust the solution is by determining how steep these
% curves are around the minimum. Use the zoom tool to visualize behavior near the
% minimum of the cost function

if plottrue
    
    figure
    plot(P(1,:),fval,'.')
    title('Moment');
    hold on
    plot(P(1,i),fval(i),'r.')
    plot(P(1,i0),fval(i0),'c.')
    plot([mopt mopt],ylim,'m--');
    hold off
    
    figure
    plot(P(2,:),fval,'.')
    title('Inclination');
    hold on
    plot(P(2,i),fval(i),'r.')
    plot(P(2,i0),fval(i0),'c.')
    plot([iopt iopt],ylim,'m--');
    hold off
    
    figure
    plot(P(3,:),fval,'.')
    title('Declination');
    hold on
    plot(P(3,i),fval(i),'r.')
    plot(P(3,i0),fval(i0),'c.')
    plot([dopt dopt],ylim,'m--');
    hold off
    
    figure
    plot(P(6,:),fval,'.')
    title('Height');
    hold on
    plot(P(6,i),fval(i),'r.')
    plot(P(6,i0),fval(i0),'c.')
    plot([hopt hopt],ylim,'m--');
    hold off
    
    figure
    plot(P(4,:),fval,'.')
    title('X displacement');
    hold on
    plot(P(4,i),fval(i),'r.')
    plot(P(4,i0),fval(i0),'c.')
    plot([xopt xopt],ylim,'m--');
    hold off
    
    figure
    plot(P(5,:),fval,'.')
    title('Y displacement');
    hold on
    plot(P(5,i),fval(i),'r.')
    plot(P(5,i0),fval(i0),'c.')
    plot([yopt yopt],ylim,'m--');
    hold off
    
end

Popt2=[Popt(1:4) 90-Popt(5) Popt(6)+90];
[resid,Bzmodel]=SourceFitDipoleBz(Popt2,Xc,Yc,scanc,0,METHOD);   % Compute map with calculated optimal parameters
residex=Bzmodel-scanc;  % Compute residuals for this solution

%Show maps of experimental data, model data, and residuals
figure
subplot(2,2,1);
imagesc(xc,yc,scanc);
axis xy, axis equal, axis tight;
caxis([-1 1]*max(abs(caxis)));
colorbar
title('Original Scan');
subplot(2,2,2);
imagesc(xc,yc,Bzmodel);
axis xy, axis equal, axis tight;
caxis([-1 1]*max(abs(caxis)));
colorbar
title('Model Scan');
subplot(2,2,4);
imagesc(xc,yc,residex);
axis xy, axis equal, axis tight;
caxis([-1 1]*max(abs(caxis)));
colorbar
title('Residuals');

%save figure to disk
saveas(gcf,[filepath '/Fit_' name '_M1_x' num2str(round(XY(1))) 'y' num2str(round(XY(2))) '.png'])

resids=sqrt(sum(sum(residex.^2))/sum(sum(scanc.^2)))    %Final residuals normalized by norm of experimental field map

%save inversion data to disk
[path,name,ext]=fileparts(INFILE);
nameext=[name ext];
nameext

if outputtrue
    fid=fopen([path '/' OutFileName],'r');
    header=(fid==-1);
    if ~header
        fclose(fid);
    end
    fid=fopen([path '/' OutFileName],'a+t');
    if header
        fprintf(fid,'File Name\tMoment\tInclination\tDeclination\tHeight\tResiduals\r\n');
    end
    %Note a 180 rotation about y axis is imposed here
    if hopt < 0
        fprintf(fid,'%s\t%1.5d\t%1.5d\t%1.5d\t%1.5d\t%1.5d\r\n',nameext,mopt,-iopt,mod(180-dopt,360),-hopt,resids);
    else
        fprintf(fid,'%s\t%1.5d\t%1.5d\t%1.5d\t%1.5d\t%1.5d\r\n',nameext,mopt,-iopt,mod(360-dopt,360),hopt,resids);
    end
    fclose(fid);
end




